# 安装依赖库
#!pip install faiss-cpu sentence-transformers numpy requests

# 导入库
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import requests

# 设置 DeepSeek API 密钥和端点
DEEPSEEK_API_KEY = "sk-xxxxx"  # 替换为你的 DeepSeek API 密钥
DEEPSEEK_API_URL = "https://api.deepseek.com/v1/chat/completions"  # 替换为 DeepSeek API 的 URL

# 文学名著知识库（《红楼梦》和《三国演义》经典段落）
documents = [
    "《红楼梦》开篇写道：“满纸荒唐言，一把辛酸泪。都云作者痴，谁解其中味？”",
    "《红楼梦》中贾宝玉说：“女儿是水做的骨肉，男人是泥做的骨肉。我见了女儿便清爽，见了男子便觉浊臭逼人。”",
    "《三国演义》开篇写道：“话说天下大势，分久必合，合久必分。”",
    "《三国演义》中诸葛亮说：“鞠躬尽瘁，死而后已。”",
    "《红楼梦》中林黛玉写道：“花谢花飞飞满天，红消香断有谁怜？”",
    "《三国演义》中曹操说：“宁教我负天下人，休教天下人负我。”",
]

# 加载嵌入模型
embedder = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')  # 支持多语言的模型

# 生成嵌入向量并构建 FAISS 索引
document_embeddings = embedder.encode(documents)
dimension = document_embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(document_embeddings)

# 检索相关文档
def retrieve_documents(query, top_k=2):
    """
    根据用户问题检索最相关的文档。
    :param query: 用户问题
    :param top_k: 返回的文档数量
    :return: 相关文档列表
    """
    query_embedding = embedder.encode([query])
    distances, indices = index.search(query_embedding, top_k)
    return [documents[i] for i in indices[0]]

# 调用 DeepSeek API 生成答案
def generate_answer(query, context):
    """
    调用 DeepSeek API 生成答案。
    :param query: 用户问题
    :param context: 检索到的上下文
    :return: 生成的答案
    """
    # 构建请求数据
    prompt = f"上下文：{context}\n\n问题：{query}\n回答："
    #print("提示词：", prompt)
    data = {
        "model": "deepseek-chat",  # 替换为 DeepSeek 的模型名称
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 100,
        "temperature": 0.7,
    }
    
    # 调用 DeepSeek API
    headers = {
        "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
        "Content-Type": "application/json",
    }
    response = requests.post(DEEPSEEK_API_URL, json=data, headers=headers)
    
    # 解析响应
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"].strip()
    else:
        raise Exception(f"DeepSeek API 调用失败: {response.status_code}, {response.text}")

# 示例运行
def run_example(query):
    """
    运行示例：检索文档并生成答案。
    :param query: 用户问题
    """
    # 检索相关文档
    retrieved_docs = retrieve_documents(query)
    print("检索到的文档：", retrieved_docs)
    
    # 生成答案
    context = "\n".join(retrieved_docs)
    answer = generate_answer(query, context)
    print("回答：", answer)

# 示例问题
query = "《红楼梦》中贾宝玉对男女的看法是什么？"
print("问题：", query)
run_example(query)

# 另一个示例问题
query = "《三国演义》开篇的第一句话是什么？"
print("\n问题：", query)
run_example(query)